# -*- coding: utf-8 -*-
import torch
import os
import datetime
from transformers.models import BertModel, BertTokenizer, BertConfig
from settings import path

import torch
import torch.nn as nn
from src import bert_train

class_list = [line.strip() for line in open(path.path_class_txt, encoding='utf-8')]

# 模型训练+预测的时候, 放开下一行代码, 在GPU上运行.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 训练设备，如果GPU可用，则为cuda，否则为cpu
num_classes = len(class_list)
num_epochs = 10  # epoch数
batch_size = 128  # mini-batch大小
pad_size = 32  # 每句话处理成的长度(短填长切)
learning_rate = 5e-5  # 学习率
bert_model = BertModel.from_pretrained(path.path_bert_base_chinese)
tokenizer = BertTokenizer.from_pretrained(path.path_bert_base_chinese)  # BERT模型的分词器
bert_config = BertConfig.from_pretrained(path.path_bert_base_chinese)  # BERT模型的配置
hidden_size = 768  # BERT模型的隐藏层大小


class BertClassifier(nn.Module):
    """
    BERT + 全连接层的分类模型。
    """
    def __init__(self):
        """
        初始化模型，包括BERT和全连接层。
        """
        super(BertClassifier, self).__init__()
        # 加载预训练的BERT模型
        self.bert = BertModel.from_pretrained(path.path_bert_model)
        # 全连接层：将BERT的隐藏状态映射到类别数
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, input_ids, attention_mask, return_hidden=False):
        # x: 模型输入，包含句子、句子长度和填充掩码。
        # _是占位符，接收模型的所有输出，而 pooled 是池化的结果,将整个句子的信息压缩成一个固定长度的向量
        _, pooled = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=False)
        # 模型输出，用于文本分类
        out = self.fc(pooled)
        if return_hidden:
            return out, pooled  # 返回logits和隐藏状态
        return out

class BiLSTMClassifier(nn.Module):
    """
    BiLSTM + 全连接层的分类模型，作为学生模型。
    """
    def __init__(self, embed_size=128, hidden_size=256, num_layers=2, num_classes=num_classes, dropout=0.3):
        """
        初始化BiLSTM模型。
        参数：
            embed_size: 嵌入维度。
            hidden_size: LSTM隐藏状态维度。
            num_layers: LSTM层数。
            num_classes: 分类类别数。
            dropout: Dropout比例。
        """
        super(BiLSTMClassifier, self).__init__()
        vocab_size = tokenizer.vocab_size  # 从BERT分词器动态获取词汇表大小
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, bidirectional=True, batch_first=True, dropout=dropout)
        self.hidden_projection = nn.Linear(hidden_size * 2, hidden_size)  # 映射到BERT隐藏状态维度
        self.fc = nn.Linear(hidden_size * 2, num_classes)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input_ids, attention_mask, return_hidden=False):
        """
        前向传播，仅在嵌入层使用 attention_mask 进行掩码处理。
        参数：
            input_ids: 输入的token ID，形状为 [batch_size, seq_len]。
            attention_mask: 注意力掩码，形状为 [batch_size, seq_len]，1 表示有效 token，0 表示填充 token。
            return_hidden: 是否返回隐藏状态。
        返回：
            logits: 分类logits，形状为 [batch_size, num_classes]。
            hidden: 最后一时间步的隐藏状态（若 return_hidden=True），形状为 [batch_size, hidden_size*2]。
        """
        # 嵌入层
        embed = self.embedding(input_ids)  # [batch_size, seq_len, embed_size]

        # 使用 attention_mask 掩码填充 token 的嵌入（核心处理）
        attention_mask = attention_mask.unsqueeze(-1)  # [batch_size, seq_len, 1]
        embed = embed * attention_mask  # 将填充 token 的嵌入置为 0

        # LSTM 层
        lstm_out, (hidden, _) = self.lstm(embed)  # lstm_out: [batch_size, seq_len, hidden_size*2]

        # 取最后一时间步的隐藏状态（填充 token 已置 0，无需再次处理）
        hidden = lstm_out[:, -1, :]  # [batch_size, hidden_size*2]

        # Dropout 和全连接层
        hidden = self.dropout(hidden)  # [batch_size, hidden_size*2]
        logits = self.fc(hidden)  # [batch_size, num_classes]

        if return_hidden:
            projected_hidden = self.hidden_projection(hidden)  # 映射到768维
            return logits, projected_hidden
        return logits


